import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from model.attack import Attack
import datetime
import csv
import kornia

from noise_layers.jpeg import Jpeg, JpegSS
from noise_layers.gaussian_noise import Gaussian_Noise
from noise_layers.quantization import Quantization
from noise_layers.salt_pepper_noise import SP
from noise_layers.colorjitter import ColorJitter

def train(model, train_loader, test_loader, num_epochs, model_dir, log_file, save_image_dir):


    mse_loss = nn.MSELoss().cuda()
    # ssim_loss = kornia.losses.MS_SSIMLoss().cuda()
    # cross_loss = nn.CrossEntropyLoss().cuda()

    # optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    model.cuda()

    """JPEG压缩"""
    jpeg = Jpeg().cuda()

    """高斯噪声"""
    # gaussian = Gaussian_Noise().cuda()

    """量化噪声"""
    # quantization = Quantization().cuda()

    """椒盐噪声"""
    # salt = SP().cuda()

    """亮度/对比度/饱和度/色相"""
    # distortion = 'Brightness'   # 亮度
    # distortion = 'Contrast'   # 对比度
    # distortion = 'Saturation'   # 饱和度
    # distortion = 'Hue'   # 色相
    # colorjitter = ColorJitter(distortion).cuda()

    mini_bce_loss = 100.0  # 跟踪最佳准确率

    with open(log_file, "w") as f:
        writer = csv.writer(f)
        writer.writerow(["Epoch", "trainLoss", "testLoss"])

        for epoch in range(num_epochs):
            is_first_save = True
            train_running_loss = 0.0
            test_running_loss = 0.0

            # 训练阶段
            model.train()
            with torch.enable_grad():
                for train_images, _ in train_loader:
                    train_images = train_images.cuda()
                    # train_images_a = train_images.clone()
                    # train_images_b = train_images.clone()
                    # noise_images = gaussian(train_images)
                    # noise_images = colorjitter(train_images)
                    noise_images = jpeg(train_images)
                    # noise_images = quantization(train_images)
                    # noise_images = salt(train_images)

                    optimizer.zero_grad()
                    attack_images = model(train_images)
                    loss = mse_loss(attack_images, noise_images)
                    # loss = ssim_loss(attack_images, noise_images)
                    # loss = cross_loss(attack_images, noise_images)

                    loss.backward()
                    optimizer.step()
                    train_running_loss += loss.item()

            train_epoch_loss = train_running_loss / len(train_loader)

            # 测试阶段
            model.eval()
            with torch.no_grad():
                for images, _ in test_loader:
                    images = images.cuda()

                    noise_images = jpeg(images)
                    # noise_images = gaussian(images)
                    # noise_images = colorjitter(images)
                    # noise_images = quantization(images)
                    # noise_images = salt(images)
                    attack_images = model(images)
                    loss = mse_loss(attack_images, noise_images)
                    # loss = ssim_loss(attack_images, noise_images)
                    # loss = cross_loss(attack_images, noise_images)
                    test_running_loss += loss.item()

                    if is_first_save:  # 只保存第一个批次的图像
                        is_first_save = False
                        save_images(images.cpu()[:8, :, :, :],
                                  noise_images[:8, :, :, :].cpu(), attack_images[:8, :, :, :].cpu(), epoch, save_image_dir, normalize)

            test_epoch_loss = test_running_loss / len(test_loader)

            # 格式化损失为保留4位小数的字符串
            train_loss_str = "{:.4f}".format(train_epoch_loss)
            test_loss_str = "{:.4f}".format(test_epoch_loss)

            writer.writerow([epoch+1, train_loss_str, test_loss_str])

            print(f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_epoch_loss:.4f}, Test Loss: {test_epoch_loss:.4f}")

            # 保存最佳模型
            if loss < mini_bce_loss:
                mini_bce_loss = loss
                save_path = os.path.join(model_dir, "best_model.pth")  # 模型保存路径
                torch.save(model.state_dict(), save_path)  # 保存当前最佳模型权重

    print(f"Training log saved to {log_file}.")
    print("Training finished.")

def save_images(original_images, noise_images, attack_images, epoch, folder, normalize):
    images = original_images[:original_images.shape[0], :, :, :].cpu()
    noise_images = noise_images[:noise_images.shape[0], :, :, :].cpu()
    attack_images = attack_images[:attack_images.shape[0], :, :, :].cpu()

    # scale values to range [0, 1] from original range of [-1, 1]
    if normalize:
        images = (images + 1) / 2
        noise_images = (noise_images + 1) / 2
        attack_images = (attack_images + 1) / 2

    revise_image = torch.abs(images - noise_images) * 5
    revise_image_ = torch.abs(images - attack_images) * 5

    stacked_images = torch.cat([images, noise_images, attack_images, revise_image, revise_image_], dim=0)
    filename = os.path.join(folder, 'epoch-{}.png'.format(epoch + 1))
    # torchvision.utils.save_image(stacked_images, filename=filename, original_images.shape[0], normalize=False)
    torchvision.utils.save_image(stacked_images, filename, normalize=False)


if __name__ == '__main__':
    model = Attack(channels=64, attack=True)

    file_dir = "attack_runs"
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)

    # 获取当前日期和时间
    current_time = datetime.datetime.now()
    # 格式化日期和时间
    time_str = current_time.strftime("%Y-%m-%d_%H-%M-%S")

    # name = "run-Gaussian(0.06)" + "_" + time_str
    name = "run-JPEG(50)" + "_" + time_str
    # name = "run-Brightness" + "_" + time_str
    # name = "run-Contrast" + "_" + time_str
    # name = "run-Saturation" + "_" + time_str
    # name = "run-Hue" + "_" + time_str
    # name = "run-Quantization" + "_" + time_str
    # name = "run-salt" + "_" + time_str
    train_dir = os.path.join(file_dir, name)
    os.makedirs(train_dir, exist_ok=True)

    train_data_dir = "data/train"  # 训练数据
    test_data_dir = "data/val"     # 测试数据
    batch_size = 64
    num_epochs = 100

    model_dir = "model"
    model_dir = os.path.join(train_dir, model_dir)
    os.makedirs(model_dir, exist_ok=True)

    log_dir = "log"
    log_dir = os.path.join(train_dir, log_dir)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    log_file = os.path.join(log_dir, "log.csv")

    save_image_dir = "image"
    save_image_dir = os.path.join(train_dir, save_image_dir)
    os.makedirs(save_image_dir, exist_ok=True)

    normalize = False
    image_size = 128
    # 定义图像预处理的转换操作
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop((image_size, image_size), pad_if_needed=True),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) if normalize else transforms.Compose([]),
        ]),
        'test': transforms.Compose([
            transforms.CenterCrop((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) if normalize else transforms.Compose([]),
        ])
    }

    # 创建ImageFolder数据集实例
    train_dataset = ImageFolder(train_data_dir, transform=data_transforms['train'])
    test_dataset = ImageFolder(test_data_dir, transform=data_transforms['test'])

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    train(model, train_loader, test_loader, num_epochs, model_dir, log_file, save_image_dir)



def differ_Jpegss_jpeg(test_loader, save_image_dir):

    device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
    jpeg_ss = JpegSS().cuda()
    jpeg = Jpeg().cuda()

    mse_loss = nn.MSELoss().cuda()
    ssim_loss = kornia.losses.MS_SSIMLoss().cuda()
    is_first_save = True
    test_running_loss = 0.0

    for images, _ in test_loader:
        images = images.cuda()
        jpeg_images = jpeg(images)
        jpegss_images = jpeg_ss([images, images])[0]

        # loss = mse_loss(jpeg_images, jpegss_images)
        loss = ssim_loss(jpeg_images, jpegss_images)
        test_running_loss += loss.item()

        # if is_first_save:  # 只保存第一个批次的图像
        #     is_first_save = False
        #     jpeg_images = jpeg_images.to(float)
        #     jpegss_images = jpegss_images.to(float)
        #     save_images(jpeg_images.cpu()[:8, :, :, :],
        #                 jpegss_images[:8, :, :, :].cpu(), 0, save_image_dir)

    test_epoch_loss = test_running_loss / len(test_loader)
    print("testLoss:{}".format(test_epoch_loss))